/*
 * Copyright 2021-2022 Open Kunlun Technology <https://www.openkunlun.io>
 */

package io.openkunlun.scaladsl.context

import akka.actor.{ ExtendedActorSystem, Extension, ExtensionId, ExtensionIdProvider }
import akka.grpc.scaladsl.Metadata
import io.openkunlun.scaladsl.context.Context.ContextTag

/**
 * @author ericxin.
 */
object AkkaGrpcContext extends ExtensionId[AkkaGrpcContext] with ExtensionIdProvider {
  override def lookup: ExtensionId[AkkaGrpcContext] = AkkaGrpcContext
  override def createExtension(system: ExtendedActorSystem) = new AkkaGrpcContext(system)
}
class AkkaGrpcContext(system: ExtendedActorSystem) extends Extension {

  private val contextPropagation: ContextStorage = ContextStorage.current(system)

  def init(): Unit = {
    contextPropagation.init()
  }

  def store[T](metadata: Metadata)(f: => T): T = {
    val tags = metadata.asList.map(it => it._1 -> metadata.getText(it._1))
      .filter(_._2.isDefined)
      .map(it => ContextTag(it._1, it._2.get))
    contextPropagation.store(new Context(tags))(f)
  }

  def get(): Context = {
    contextPropagation.get()
  }

  def tags(): Seq[ContextTag] = {
    contextPropagation.get().tags
  }
}
